from easy_io import read_mat_file, H5Writer
import numpy as np
import h5py
from scipy.ndimage import zoom


def read_partly_from_h5_file(h5_file_obj):
    def helper(path_starts_shape):
        path, starts, shape = path_starts_shape
        slices = tuple(slice(max(st, 0), st + sp) for st, sp in zip(starts, shape))
        raw = h5_file_obj[path][slices]
        mshape = h5_file_obj[path].shape
        pad_width = np.asarray([(0 - st, st + sp - msp) for st, sp, msp in zip(starts, shape, mshape)], 'int')
        pad_width[pad_width < 0] = 0
        paded = np.pad(raw, pad_width, 'constant', constant_values=-1024)
        assert paded.shape == tuple(shape)
        return paded

    return helper


def gen(candidate_mat_file, data_h5_file, min_edge_length):
    candidate_arr = read_mat_file(candidate_mat_file)
    candidate_lst = []
    for s in candidate_arr:
        scanid = str(s[0][0])
        assert s[1].ndim == 2
        for c in s[1]:
            c = np.asarray(c[:4], 'int')
            candidate_lst.append((scanid, c[:3], c[3]))

    with h5py.File(data_h5_file, 'r') as f:
        reader = read_partly_from_h5_file(f)
        for i, (scanid, midpoint, diameter) in enumerate(candidate_lst):
            # midpoint = np.array([midpoint[1], midpoint[0], midpoint[2]], 'int')
            if scanid in f:
                spacing = np.asarray(f[scanid]['vol'].attrs['spacing'])
                assert spacing[0] == spacing[1], scanid

                unit = np.min(spacing)
                edge_length = max(min_edge_length, diameter)
                crop_shape = np.asarray(np.ceil((edge_length - 1) * unit / spacing + 1), 'int')
                starts = np.asarray(midpoint - crop_shape // 2, 'int')
                vol = reader((scanid + '/vol', starts, crop_shape))

                vol = zoom(vol, min_edge_length / np.asarray(vol.shape), output='float32', order=1, mode='nearest')
                assert np.all(np.asarray(vol.shape) == min_edge_length)
                yield str(i), vol
            else:
                print(scanid)


if __name__ == '__main__':
    H5Writer(
        '/ssd_1t/Kaggle/lidc_candidate_vol.hdf5',
        'w',
        gen(
            candidate_mat_file='/data_4t/Kaggle/candidates/lidc_candis.mat',
            data_h5_file='/data_4t/Kaggle/lidc&kaggle/train_data.hdf5',
            min_edge_length=68,
        )
    )
